// %%%%%%%%%%%%%%%%%%%%%%% HILO_DEBUGGING.H %%%%%%%%%%%%%%%%%%%%%%%
// Contains some functionality related to debugging.
// ----------------------------------------------------------------

#include <Windows.h>
#include <Python.h>

#include "globals.h"
#include "hilo_HANDLE.h"
#include "hilo_CONTEXT.h"
#include "procutil.h"
#pragma once


//
// Small buffer for the INT3 breakpoint (only one byte needed to write; formatted for nicety.
//

static BYTE BP_INT3[] = {0xCC, 0x00, 0x00, 0x00};

//
// Main debugging thread to be used for debugging a process.
// Handles setup and initialization for debugging
//	Requires a pointer to a hilo_DebugInfo object (on the heap) passed as the paramater.
//	Also requires the debugging start address
//

DWORD WINAPI hilo_DebugThread( LPVOID lpParam );

//
// Contains information that is passed onto and stored as state in a debugging thread.
// Ideally we wish to make our implementation general enough so that multiple places may be debugged
//		but due to time constraints, we are only limiting ourself to 1 BP/trace at a time.
//		An design to implement many breakpoints is to pass in a Process object and provide some way of registering
//			breakpoints with callbacks, and then performing the function calls as they happen.
//

class hilo_DebugInfo {
public:
	//
	// A hilo_HANDLE object (process handle) to debug
	//

	hilo_HANDLE		*_hProc;

	//
	// The callback function to use when a BP is hit.
	// The callback function accepts the instruction address and a context object.
	// Modifications made to the context object are reflected in the thread continuation.
	//

	PyObject		*pCallback;

	//
	// The address of the breakpoint set by this debugging instance.
	//

	DWORD			dwAddress;

	//
	// The byte written over to place the INT3 breakpoint.
	//

	DWORD			dwByteValue;

	//
	// Instruction counter for the trace.
	//	(Because Python is stupid)
	//

	DWORD			dwInstrCounter;

	// 
	// A HANDLE to the synchronizing semaphore.
	// This semaphore is signaled when the debugging thread is ready for a Call.
	// The semaphore is automatically closed when the destructor is called.
	//

	HANDLE			hReadyState;

	//
	// Constructor for the hilo_DebugInfo object.
	// Accepts a Python HANDLE object, a callback object, and the address to start debugging.
	// Does not actually place or read from memory; merely sets up the object.
	//

	__inline hilo_DebugInfo( hilo_HANDLE *h, PyObject *callback, DWORD _dwAddress ) {
		Py_XINCREF( h );
		Py_XINCREF( callback );
		this->_hProc = h;
		this->pCallback = callback;
		this->dwAddress = _dwAddress;
		this->hReadyState = CreateSemaphore( NULL, 0, 1, NULL );
		this->dwInstrCounter = 1;
	}

	//
	// Destructor for the hilo_DebugInfo object.
	// Takes care to decrement references to Python objects set earlier in constructor.
	//

	__inline ~hilo_DebugInfo( void ) {
		Py_XDECREF( this->_hProc );
		Py_XDECREF( this->pCallback );
		CloseHandle( this->hReadyState );
	}

	//
	// Places a breakpoint at the given address.
	// Saves the values written over before placing.
	// Returns zero only if an error occurred.
	//

	__inline BOOL __fastcall placeBreakpoint( DWORD _dwAddress ) {
		this->dwAddress = _dwAddress;	// set address
		BOOL success = readMemory( this->_hProc->hHandle, _dwAddress, 1, &(this->dwByteValue) ); // read from memory

		if ( success == NULL ) { // report failure
			DEBUG_PRINT( "Critical memory read failed at 0x%08X", _dwAddress );
		} else { // success read
			success = writeMemory( this->_hProc->hHandle, _dwAddress, 1, BP_INT3 ); // write the breakpoint
			// flush instruction cache
			FlushInstructionCache( this->_hProc->hHandle, (LPCVOID)this->dwAddress, 20 );
			if ( success == NULL ) {
				DEBUG_PRINT( "BP write failed at 0x%08X", _dwAddress );
			}
		}

		return success;
	}

	//
	// Restores the instruction at the breakpoint address.
	// Returns zero only if an error occurred.
	//

	__inline BOOL restoreInstruction( DWORD eip ) {
		BOOL success = writeMemory( this->_hProc->hHandle, this->dwAddress, 1, &(this->dwByteValue) );
		DEBUG_PRINT( "restoring instr" );
		if ( success = NULL ) { // report failure
			DEBUG_PRINT( "Instruction restore failed at 0x%08X", this->dwAddress );
		} else if ( eip != this->dwAddress ) {
			BOOL success = writeMemory( this->_hProc->hHandle, eip, 1, &(this->dwByteValue) );
			DEBUG_PRINT( "restoring secondary instr" );
			if ( success = NULL ) { // report failure
				DEBUG_PRINT( "Secondary instruction restore failed at 0x%08X", eip );
			}
		}

		// flush instruction cache
		FlushInstructionCache( this->_hProc->hHandle, (LPCVOID)this->dwAddress, 20 );

		return success;
	}
	
	//
	// Performs the Python callback given a pointer to the CONTEXT structure to use.
	// Note: the CONTEXT structure is modified by the callback.
	// Returns zero only on error or a quit signal (synonymous).
	// Otherwise, returns the next address to place an instruction at.
	//

	__inline DWORD __fastcall performCallback( DEBUG_EVENT const * de, CONTEXT * cContext ) {
		DWORD nextAddress = 0;		// return value (next address to place BP at)
		PyObject *args, *res;		// arguments to pass to callback
		hilo_CONTEXT *hiloContext;	// Python context object

		// create the context
		hiloContext = (hilo_CONTEXT*)hilo_CINT::fromContext( cContext );
		Py_XINCREF( hiloContext );

		DEBUG_PRINT("Performing callback");
		DEBUG_PRINT_VAR( de->u.Exception.ExceptionRecord.ExceptionAddress );
		// form arguments and make the call (then decrement counter)
		args = Py_BuildValue( "(IOI)", de->u.Exception.ExceptionRecord.ExceptionAddress, hiloContext, this->dwInstrCounter );
		Py_XINCREF( args );
		DEBUG_PRINT("Finished forming args");
		res = PyObject_CallObject( this->pCallback, args );
		Py_XDECREF( args );

		DEBUG_PRINT("Finished performing callback");
		DEBUG_PRINT_VAR( PyInt_AsUnsignedLongMask(res) );

		// on success, modify context and set the return offset
		if ( res != NULL ) {
			*cContext = hiloContext->cContext;
			nextAddress = PyInt_AsUnsignedLongMask(res);
			DEBUG_PRINT_VAR( nextAddress );
			Py_DECREF( res );
		} else {
			DEBUG_PRINT( "Error produced in callback" );
		}

		Py_XDECREF( hiloContext );

		return nextAddress;
	}

	//
	// Creates the debugging thread given this dbgEvent info's information.
	// Returns a HANDLE to the thread given, or NULL on failure.
	//
	__inline HANDLE CreateDebuggingThread( void ) {
		return CreateThread( NULL, 1024, hilo_DebugThread, this, NULL, NULL );
	}
};